from typing import List
from llama_index.core.schema import BaseNode, NodeWithScore
from llama_index.core import QueryBundle # 修正了导入路径
# 修正了导入路径，以兼容新版本的 LlamaIndex
from llama_index.core.postprocessor.sbert_rerank import SentenceTransformerRerank

from .node_postprocessor_interface import NodePostprocessorInterface


class SentenceTransformerRerankProcessor(NodePostprocessorInterface):
    """
    一个实现了 NodePostprocessorInterface 的具体后处理器，
    它在内部使用 LlamaIndex 的 SentenceTransformerRerank 来对节点进行重排序。
    """

    def __init__(self, model_name: str, top_n: int = 3):
        """
        初始化重排序器。

        Args:
            model_name: 要使用的 Cross-Encoder 模型名称。
                        例如: 'BAAI/bge-reranker-base'
            top_n: 重排序后要返回的前 N 个结果。
        """
        # 在内部创建一个 LlamaIndex Reranker 的实例
        self._reranker = SentenceTransformerRerank(
            model=model_name,
            top_n=top_n,
        )

    async def postprocess_nodes(
        self, nodes: List[NodeWithScore], query_bundle: QueryBundle, **kwargs
    ) -> List[NodeWithScore]:
        """
        使用内部的 SentenceTransformerRerank 实例来执行重排序。
        """
        # LlamaIndex 的 postprocessor 需要一个包含原始查询字符串的 QueryBundle
        return self._reranker.postprocess_nodes(nodes, query_bundle=query_bundle) 